import pickle
import multiprocessing as mp
import trimesh
from autolab_core import RigidTransform
from copy import deepcopy
import numpy as np 
import json
import os
from grasping.appoach_grasp  import ApproachParallelJawPtGrasp3D
import torch
from nms import nms


def get_grasp_and_score(grasps):
    len_grasp = len(grasps)
    score_arr = np.zeros((len_grasp, 1), dtype=np.float32)
    grasp_arr = np.zeros((len_grasp, 9), dtype=np.float32)

    for i, g in enumerate(grasps):
        score_arr[i] = g.grasp_score.sum()
        grasp_arr[i][:3] = g.center
        grasp_arr[i][3:6] = g.approached_full_axis[:, 0]
        grasp_arr[i][6:9] = g.axis
    return score_arr, grasp_arr


def prune_grasp_method(dataset_dict, key):
        path = dataset_dict[key]['path']
        mesh = trimesh.load_mesh(path)
        files = os.listdir('/home/v-wewei/code/two_stage_pointnet/generated_grasp_10000/{}'.format(key))

        for file_name in files:
            if os.path.splitext(file_name)[0].startswith('first'):
                with open('/home/v-wewei/code/two_stage_pointnet/generated_grasp_10000/{}/{}'.format(key, file_name), 'rb') as f:
                    grasps = pickle.load(f)
                    if use_gpu:
                        score_arr, grasp_arr = get_grasp_and_score(grasps)
                        score_tensor = torch.from_numpy(score_arr).view(-1).cuda()
                        grasp_tensor = torch.from_numpy(grasp_arr).cuda()
                        keep, num_to_keep, parent_object_index = nms(grasp_tensor, score_tensor, metric=.005, top_k=20000)
                        keep_grasps = keep[:num_to_keep]
                        keep_list = keep_grasps.tolist()
                        pruned_grasps = [grasps[idx] for idx in keep_list]
                    print(key, 'prune_grasps length is : ', len(pruned_grasps))
                    with open('/home/v-wewei/code/two_stage_pointnet/generated_grasp_10000/{}/final_{}.pickle'.format(key, len(pruned_grasps)), 'wb') as f_1:
                        pickle.dump(pruned_grasps, f_1)
        print('finish')

if __name__ == "__main__":
    use_gpu = True
    flag = False
    # pool = mp.Pool(processes=3)
    with open('/home/v-wewei/finish_stl_02/dataset_dict.json', 'r') as f:
        dataset_dict = json.load(fp=f)
        for key in dataset_dict.keys():
            files = os.listdir('/home/v-wewei/code/two_stage_pointnet/generated_grasp_10000/{}'.format(key))
            # print(files)
            if len(files) == 3:
                continue

            # if key.startswith('011_banana'):
            #     continue
            # if key.startswith('043'):
            #     continue
            print(key)

            # pool.apply_async(prune_grasp_method,args=(dataset_dict,key,))
            prune_grasp_method(dataset_dict, key)


        # pool.close()
        # pool.join()